#include "./beeman.h"

static void mds_beeman_predict( mds_frame *frame,  const mds_beeman_data *data )
{
	assert( frame );
	assert( data );
	assert( frame->sets );
	assert( frame->num_sets >= MDS_NUM_SETS_FOR_BEEMAN );
	assert( frame->num_particles > 0 );
	assert( frame->pbc_for_r );
	{
		mds_vector       * position   = frame->sets[MDS_POSITION];
		const mds_vector * velocity   = frame->sets[MDS_VELOCITY];
		const mds_vector * curr_acc   = frame->sets[MDS_ACCELERATION];
		mds_vector       * prev_acc   = frame->sets[MDS_PREV_ACCELER];
		mds_vector       * aged_acc   = frame->sets[MDS_AGED_ACCELER];
		const size_t       n          = frame->num_particles;
		const Real         dt         = data->dt;
		const Real         sixth_dtsq = data->sixth_dtsq;
		mds_pbc_proc       pbc        = frame->pbc_for_r;
		const mds_vector  *box        = & frame->box;
		mds_vector         four_a0;
		mds_vector         dr;
		register size_t    i;
		
		
		for( i=0; i < n; ++i )
		{
			mds_vector       *r  = &position[i];
			const mds_vector *a0 = &curr_acc[i];
			mds_vector       *a1 = &prev_acc[i];
			mds_vector       *a2 = &aged_acc[i];
			
			mds_vector_mulset( &four_a0, MDS_F(4.0), a0 );
			mds_vector_mulvec( &dr, sixth_dtsq, a1, &four_a0 );
			mds_vector_muladd( &dr, dt, &velocity[i] );
			mds_vector_add( r, &dr );
			
			pbc( r, box );
			
			mds_vector_cpy(a2,a1); /* a2 <- a1 */
			mds_vector_cpy(a1,a0); /* a1 <- a0 */
			/* a0 is ready to receive new accelerations */
		}
	}
	
}

static void mds_beeman_correct( mds_frame *frame,  const mds_beeman_data *data )
{
	assert( frame );
	assert( data );
	assert( frame->sets );
	assert( frame->num_sets >= MDS_NUM_SETS_FOR_BEEMAN);
	assert( frame->num_particles > 0 );
	assert( frame->pbc_for_r );
	{
		/* mds_vector       * position   = frame->sets[MDS_POSITION]; */
		mds_vector       * velocity   = frame->sets[MDS_VELOCITY];
		const mds_vector * curr_acc   = frame->sets[MDS_ACCELERATION];
		mds_vector       * prev_acc   = frame->sets[MDS_PREV_ACCELER];
		mds_vector       * aged_acc   = frame->sets[MDS_AGED_ACCELER];
		const size_t       n          = frame->num_particles;
		const Real         sixth_dt   = data->sixth_dt;
		mds_vector         da;
		register size_t    i;
		Real               K          = 0;
		
		for( i=0; i < n ; ++i )
		{
			mds_vector       *v  = &velocity[i];
			const mds_vector *a0 = &curr_acc[i];
			mds_vector       *a1 = &prev_acc[i];
			mds_vector       *a2 = &aged_acc[i];
			
			mds_vector_mulset( &da, MDS_F(2.0), a0 );
			mds_vector_muladd( &da, MDS_F(5.0), a1 );
			mds_vector_sub(    &da,    a2 );
			mds_vector_muladd( v, sixth_dt, &da );
			
			
			K += (v->info.mass) * mds_vector_normsq( v );
			
		}
		
		frame->K = MDS_F(0.5) * K;
	}
}



static void mds_beeman_init_dt( mds_frame *frame, Real dt, mds_beeman_data *data )
{
	assert( frame );
	assert( data );
	assert( frame->sets );
	assert( frame->num_sets >= 5 );
	assert( frame->num_particles > 0 );
	fprintf( stderr, "mds_beeman_init dt=%g\n", dt );
	data->dt         = dt;
	data->sixth_dt   = dt/MDS_F(6.0);
	data->sixth_dtsq = data->sixth_dt * dt;
	{
		const mds_vector * curr_acc   = frame->sets[MDS_ACCELERATION];
		mds_vector       * prev_acc   = frame->sets[MDS_PREV_ACCELER];
		mds_vector       * aged_acc   = frame->sets[MDS_AGED_ACCELER];
		const size_t       n          = frame->num_particles;
		register size_t    i;
		
		for( i=0; i < n ; ++i )
		{
			const mds_vector *a0 = &curr_acc[i];
			mds_vector       *a1 = &prev_acc[i];
			mds_vector       *a2 = &aged_acc[i];
			mds_vector_ldz( a2 );
			mds_vector_cpy( a1, a0 );
		}
		
	}
}


void mds_beeman_step_load( mds_step *step, mds_beeman_data *data )
{
	assert( step );
	assert( data );
	
	step->data     = data;
	step->predict  = (mds_step_upgrade_proc) mds_beeman_predict;
	step->correct  = (mds_step_upgrade_proc) mds_beeman_correct;
	step->init_dt  = (mds_step_init_dt_proc) mds_beeman_init_dt;
	step->num_sets = MDS_NUM_SETS_FOR_BEEMAN;
}

